import os
from openai import OpenAI

class YiModel:
    API_BASE = "https://api.lingyiwanwu.com/v1"
    API_KEY = os.environ.get("YI_API_KEY")

    def __init__(self, model="yi-lightning"):  # yi-medium, yi-large-turbo
        self.model = model

    def chat(self, prompt):
        messages = [{'role': 'user', 'content': prompt}]
        client = OpenAI(
            api_key=self.API_KEY,
            base_url=self.API_BASE
        )
        try:
            completion = client.chat.completions.create(
                model=self.model,
                messages=messages
            )
            return completion.choices[0].message.content
        except Exception as e:
            print('YiModel exception:', str(e))
            return None


if __name__ == "__main__":
    prompt = """按行判断下面每行文字是什么语言。然后将它翻译成英语和简体中文。最后按csv格式输出结果：第一列为原文，第二列为2字母的语言代码，第三列为英语结果，第四列为中文结果。如果原文或译文包含逗号或者引号，按照csv的标准方式进行转义处理。结果只需要输出一遍。不要输出csv表头。如果文字无法翻译，就将这一行舍弃。
cámara mini llavero
magic john iphone 14 pro max
aceite de batana organico
hana ginseng
making the team
gus gus cenicienta
batería arrancador
best budget camera
wd my book 18tb
kit de skincare piel grasa
ridiculous inflatable swan thing
expansión series x
impermeable perro lluvia
cavitador radiofrecuencia
alfombras tapetes
let me fuxking cry book
orgasmo gel
aceite de batana organico
charms pandor
esterilizador avent philips
    """
    print(YiModel().chat(prompt))
